"""
DQN
"""
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd


EPSILON = 0.05


class Net(nn.Module):
    def __init__(self, n_states, n_actions, neuron_number=30):
        super(Net, self).__init__()

        self.fc1 = nn.Linear(n_states, neuron_number)
        self.fc1.weight.data.normal_(0, 0.1)   # initialization
        self.out = nn.Linear(neuron_number, n_actions)
        self.out.weight.data.normal_(0, 0.1)   # initialization

    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)

        actions_value = self.out(x)
        return actions_value

        # return self.network(x)


class DQN(object):
    def __init__(
            self,
            env,
            seed=None,
            neuron_number=30,
            gamma=0.99,
            lr=0.01,
            memory_capacity=2000,
            target_replace_iter=500,
            train_frequency=10,
            batch_size=128,
            net=None,
            is_save_memory_to_csv=False,
            clamp_grad=0,       # 是否剪裁梯度
            # is_save_model_params=False,  # 这个还是在main函数里自己定义比较好
    ):
        # from rl_envs.Car2dEnv import Car2dEnv
        # env: Car2dEnv
        self.env = env

        self.n_actions = env.action_space.n
        self.n_states = env.observation_space.shape[0]
        self.env_a_shape = 0 if isinstance(env.action_space.sample(), int) else env.action_space.sample().shape

        self.memory_capacity = memory_capacity
        self.net = net if net else Net

        self.gamma = gamma
        self.lr = lr
        self.target_replace_iter = target_replace_iter
        self.batch_size = batch_size
        self.train_frequency = train_frequency
        self.clamp_grad = clamp_grad

        self.eval_net, self.target_net = self.net(self.n_states, self.n_actions, neuron_number), self.net(self.n_states, self.n_actions, neuron_number)

        self.learn_step_counter = 0                                     # for target updating
        self.memory_counter = 0                                         # for storing memory
        self.memory = np.zeros((memory_capacity, self.n_states * 2 + 3))     # initialize memory
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=self.lr)
        self.loss_func = nn.MSELoss()

        self.default_dir = 'tempdir/dqn'
        os.makedirs(self.default_dir, exist_ok=True)

        self.is_save_memory_to_csv = is_save_memory_to_csv
        self.default_path = os.path.join(self.default_dir, 'dqn_memory.csv')

        self.model_save_path = os.path.join(self.default_dir, 'target_net.pth')
        self.model_params_save_path = os.path.join(self.default_dir, 'state_dict_model.pth')

        self.is_need_first_save_memory = True

        self.set_seed(seed)

    def set_seed(self, seed):
        if seed is not None:
            np.random.seed(seed)
            torch.manual_seed(seed)
            self.env.seed(seed)
            self.env.action_space.seed(seed)
            self.env.observation_space.seed(seed)

    def choose_action(self, x, epsilon=None):
        if epsilon is None:
            epsilon = EPSILON
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        # input only one sample
        rd = np.random.random()
        if rd > epsilon:
            # action = self.env.action_space.sample()
            # print(f'\n\n~~~~~~~~~~~ rd[{rd}] > epsilon[{epsilon}]!')
            actions_value = self.eval_net.forward(x)
            action = torch.max(actions_value, 1)[1].data.numpy()
            action = action[0] if self.env_a_shape == 0 else action.reshape(self.env_a_shape)  # return the argmax index
        else:   # random
            action = np.random.randint(0, self.n_actions)
            action = action if self.env_a_shape == 0 else action.reshape(self.env_a_shape)
        return action

    def store_transition(self, s, a, r, done, s_):
        transition = np.hstack((s, [a, r, done], s_))
        # replace the old memory with new memory
        index = self.memory_counter % self.memory_capacity
        self.memory[index, :] = transition
        self.memory_counter += 1

    def save_memory_to_csv(self, path=None):
        if path is None:
            path = self.default_path

        df = pd.DataFrame(self.memory)

        # s, a, r, done, s_2
        columns = []
        current_s_column_name = [f'curr_s_{i}' for i in range(self.n_states)]
        mid_column_name = ['a', 'r', 'done']
        next_s_column_name = [f'next_s_{i}' for i in range(self.n_states)]
        columns.extend(current_s_column_name)
        columns.extend(mid_column_name)
        columns.extend(next_s_column_name)

        df.columns = columns
        df.to_csv(path, index=False)
        print(f'\n\n=== dqn.save_memory_to_csv: save memory to path`{path}`!')

        self.is_need_first_save_memory = False

    def load_memory_from_csv(self, path=None):
        MEMORY_CAPACITY = self.memory_capacity

        if path is None:
            path = self.default_path

        # assert os.path.exists(path), f'*** dqn.load_memory_from_csv: csv_path[{path}] is not exists!'
        if not os.path.exists(path):
            return

        print(f"=== dqn.load_memory_from_csv: load memory from csv_path[{path}]!")
        df = pd.read_csv(path)
        self.memory = df.to_numpy()
        self.memory_counter = self.memory.shape[0]

        if MEMORY_CAPACITY < self.memory_counter:
            sample_index = np.random.choice(self.memory_counter, MEMORY_CAPACITY)
            self.memory = self.memory[sample_index, :]
            self.memory_counter = MEMORY_CAPACITY

            self.is_need_first_save_memory = False

    def learn(self):
        N_STATES = self.n_states
        # train_frequency = self.train_frequency

        if self.is_save_memory_to_csv and self.is_need_first_save_memory:
            self.save_memory_to_csv()

        # target parameter update
        if self.learn_step_counter % self.target_replace_iter == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())

        self.learn_step_counter += 1
        if self.learn_step_counter % self.train_frequency != 0:  # 学习频率
            return

        # sample batch transitions
        sample_index = np.random.choice(self.memory_capacity, self.batch_size)
        b_memory = self.memory[sample_index, :]
        b_s = torch.FloatTensor(b_memory[:, :N_STATES])
        b_a = torch.LongTensor(b_memory[:, N_STATES:N_STATES+1].astype(int))
        b_r = torch.FloatTensor(b_memory[:, N_STATES+1:N_STATES+2])
        b_done = torch.FloatTensor(b_memory[:, N_STATES+2:N_STATES+3])
        b_s_ = torch.FloatTensor(b_memory[:, -N_STATES:])

        # q_eval w.r.t the action in experience
        q_eval = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)

        with torch.no_grad():
            q_next = self.target_net(b_s_).detach()     # detach from graph, don't backpropagate
            q_target = b_r + self.gamma * q_next.max(1)[0].view(self.batch_size, 1)   # shape (batch, 1)

        loss = self.loss_func(q_eval, q_target)

        self.optimizer.zero_grad()
        loss.backward()
        if self.clamp_grad:
            for param in self.eval_net.parameters():
                param.grad.data.clamp_(-self.clamp_grad, self.clamp_grad)
        self.optimizer.step()


def linear_schedule(start_e: float, end_e: float, duration: int, t: int):
    slope = (end_e - start_e) / duration
    return max(slope * t + start_e, end_e)


